# 开发者：郭同源
# 开发时间：2021/10/30 13:23
import torch
import torchvision
from torch import nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(16, 22, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(22, 32, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 42, 3),
            nn.ReLU(),
            nn.Conv2d(42, 64, 3),
            nn.ReLU()
        )
        self.out_layer = nn.Sequential(
            nn.Linear(64 * 19 * 19, 25),
        )

    def forward(self, x):
        out = self.layers(x).reshape(-1, 64 * 19 * 19)
        out = self.out_layer(out)
        label = out[:, 0]
        position = out[:, 1:5]
        sort = out[:, 5:]
        return label, position, sort



if __name__ == '__main__':
    model = MyModel()
    x = torch.randn((3,3,300,300))
    output = model(x)
    # print(output[1]*300)
    print(output[1][0])
